#!/usr/bin/env python
# -*- encoding: utf-8 -*-

from elasticsearch import Elasticsearch
import click
import logging
from six import iteritems


logging.getLogger().setLevel(logging.INFO)
projects = ['CVE_SCAP', 'CVE', 'ExploitDB', 'OSVDB', 'ZDI', 'Seebug', 'wooyun_vuls', 'cxsecurity', 'JVN',
            'packetstormsecurity', 'securitytracker', 'CNVD']
project_urls = ['scap_url', 'cve_url', 'edb_url', 'osvdb_url', 'zdi_url', 'seebug_url', 'wy_url', 'cxs_url',
               'jvn_url', 'pss_url', 'st_url', 'cnvd_url']
exclude_keys = ['display_vul_id', 'title', 'date', 'opentime', 'cve_id', 'CVSS', 'description', 'content', 'CPE', 'poc', 'exp', 'url']
count = 0  # 记录库内数据总条数，用于vul_id编号


@click.command()
@click.option('--host', help='elasticsearch host address', default='http://127.0.0.1:9200/',  required=1)
def main(host):
    centrolize(host)


def result_formater(results):
    common_fields = None
    for result in results:
        result.setdefault('result', None)
        if isinstance(result['result'], dict):
            if common_fields is None:
                common_fields = set(result['result'].keys())
            else:
                common_fields &= set(result['result'].keys())
        else:
            common_fields = set()
    for result in results:
        result['result_formated'] = {}
        if not common_fields:
            result['others'] = result['result']
        elif not isinstance(result['result'], dict):
            result['others'] = result['result']
        else:
            result_formated = {}
            others = {}
            for key, value in iteritems(result['result']):
                if key in common_fields:
                    result_formated[key] = value
                else:
                    others[key] = value
            result['result_formated'] = result_formated
            result['others'] = others
    return common_fields or set(), results


def count_centdb(es):
    return es.count(index="centdb", doc_type="cent",
                         body={"query": {"match_all": {}}}
                         ).get('count', 0)


def select(es, project, fields=None, offset=0, limit=0):
    for record in es.search(index="resultdb", doc_type="result",
                            body={'query': {'bool': {'must_not': {'term': {'cent_flag': 1}},
                                                     'must': {'term': {'project': project}}}}},
                            _source_include=fields or [],
                            from_=offset, size=limit).get('hits', {}).get('hits', []):
        yield record['_source']


def select_cveid_from_centdb(es, cve_id, fields=None, offset=0):
    for record in es.search(index="centdb", doc_type="cent",
                            body={'query': {'match_phrase': {'cve_id': cve_id}}},
                            _source_include=fields or [],
                            from_=offset).get('hits', {}).get('hits', []):
        yield record['_source']


def modify_cent_flag(es, resultid, cent_flag):
    ret = es.update(index="resultdb", doc_type="result",
                    body={'doc': {'cent_flag': cent_flag}}, id=resultid)
    return ret


def insert_centdb_all(es, resultid, project, vul_id, title, date, opentime, cve_id, misc,
                  CVSS, description, content, CPE, poc, exp, VDB, log):
    obj = {
        'resultid': resultid,
        'project': project,
        'vul_id': vul_id,
        'title': title,
        'date': date,
        'opentime': opentime,
        'cve_id': cve_id,
        'misc': misc,
        'CVSS': CVSS,
        'description': description,
        'content': content,
        'CPE': CPE,
        'poc': poc,
        'exp': exp,
        'VDB': VDB,
        'log': log,
    }
    return es.index(index="centdb", doc_type="cent",
                    body=obj, id=resultid, refresh=True)


def modify_centdb_VDB(es, resultid, title, content, VDB):
    ret = es.update(index="centdb", doc_type="cent",
                    body={'doc': {'title': title, 'content': content, 'VDB': VDB}}, id=resultid, refresh=True)
    return ret


def centrolize(host):
    global projects
    global count
    resultes = Elasticsearch(hosts=host)
    centes = Elasticsearch(hosts=host)
    centes.indices.create(index='centdb', ignore=400)
    if not centes.indices.get_mapping(index='centdb', doc_type='cent'):
        centes.indices.put_mapping(index='centdb', doc_type='cent', body={
            "_all": {"enabled": True},
            "properties": {
                "project": {"type": "string", "index": "not_analyzed"},
                "resultid": {"type": "string", "index": "not_analyzed"},
                "cve_id": {"type": "string", "index": "not_analyzed"},
            }
        })
    try:
        count = count_centdb(centes)
    except:
        if not count:
            logging.info("count is 0 from centdb")
    logging.info("centdb has %d data totally", count)
    for project in projects:  # 要求各project顺序执行
        offset = 0
        limit = 20
        while True:
            results = list(select(resultes, project=project, offset=offset, limit=limit))
            nmatched = len(results)
            logging.info("loop: keyword %s get %d results", project, nmatched)
            if nmatched == 0:
                break
            offset += nmatched
            common_fields, _results = result_formater(results)
            for result in results:
                if not result:
                    continue
                project = result['project']
                taskid = result['taskid']
                resultid = '%s:%s' % (project, taskid)
                result_dict = result.get('result', {})
                date = result_dict.get('date', '2016-08-26')
                opentime = result_dict.get('opentime', '')
                cve_id = result_dict.get('cve_id', '')
                if cve_id:
                    if len(cve_id.split('-')) < 3:
                        continue
                    else:
                        vul_id = 'GASS-'+cve_id.split('-')[1]+'-00'+cve_id.split('-')[2]
                else:
                    if count < 010000:
                        vul_id = 'GASS-'+date.split('-')[0]+'-01'+'%04d'%count
                    else:
                        vul_id = 'GASS-'+date.split('-')[0]+'-%06d'%count
                CVSS = result_dict.get('CVSS', None)
                description = result_dict.get('description', None)
                CPE = result_dict.get('CPE', None)
                poc = result_dict.get('poc', None)
                exp = result_dict.get('exp', None)
                url = result_dict.get('url', None)
                VDB = {}  # dict, 用于存储具有相同cve_id的url的key-value对
                misc = {}  # dict, 用于存储各漏洞库与centdb字段（不含misc）交集外字段
                for field in common_fields:
                    if field in exclude_keys:
                        continue
                    else:
                        misc[field] = result_dict.get(field, '')
                result_ori = []
                if cve_id:
                    result_ori = list(select_cveid_from_centdb(centes, cve_id))
                if not result_ori:
                    title = [result_dict.get('title', None)]
                    content = [result_dict.get('content', None)]
                    VDB[project_urls[projects.index(project)]] = [url]
                    insert_centdb_all(centes, resultid, project, vul_id, title, date, opentime, cve_id, misc,
                                      CVSS, description, content, CPE, poc, exp, VDB, log='')
                    count += 1
                else:
                    resultid_ori = result_ori[0].get('resultid', None)
                    title_ori = result_ori[0].get('title', [])
                    title_ori.append(result_dict.get('title', None))
                    content_ori = result_ori[0].get('content', [])
                    content_ori.append(result_dict.get('content', None))
                    VDB_ori = result_ori[0].get('VDB', [])
                    if VDB_ori.get(project_urls[projects.index(project)], None):
                        VDB_ori[project_urls[projects.index(project)]].append(url)
                    else:
                        VDB_ori[project_urls[projects.index(project)]] = [url]
                    modify_centdb_VDB(centes, resultid_ori, title_ori, content_ori, VDB_ori)
                modify_cent_flag(resultes, resultid, 1)
        logging.info("keyword %s match total %d results" %(project, count))

if __name__ == '__main__':
    main()
